import types
import inspect
import torch
from torch import nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        
    def forward(self, x):
        print("model forward")
        return x
    
    def fit(self, x, y):
        print("model fit")
        logits = self.forward(x)
        return nn.functional.cross_entropy(logits, y), x
    

class Wrapper(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
    
    def forward(self, x):
        print("wrapper forward")
        return x
    

if __name__ == "__main__":
    X = torch.randn(10, 3)
    y = torch.randint(0, 2, (10,))
    model = Model()
    model = nn.DataParallel(model)
    if hasattr(model.module, "fit"):
        model.fit = types.MethodType(model.module.__class__.fit, model)
        
    # # model(X)
    model.fit(X, y)
